import { RoleType, RoleMessageType } from "../enum";
import { TextPrompt, CodePrompt } from "../prompt";

export interface BaseMessageConfig {
    /** 角色名称 */
    roleName: string;
    /** 角色类型 */
    roleType: RoleType;
    /** 角色消息类型 */
    roleMessageType: RoleMessageType;
    /** 消息内容 */
    content: string;
    /** 元数据 */
    metaDict?: Record<string, any>;
}

export class BaseMessage {
    public roleName: string;
    public roleType: RoleType;
    public roleMessageType: RoleMessageType;
    public content: string;
    public metaDict?: Record<string, string>;

    constructor(config: BaseMessageConfig) {
        this.roleName = config.roleName;
        this.roleType = config.roleType;
        this.roleMessageType = config.roleMessageType;
        this.content = config.content;
        this.metaDict = config.metaDict;
    }

    createNewInstance(content: string): BaseMessage {
        return new BaseMessage({
            roleName: this.roleName,
            roleType: this.roleType,
            roleMessageType: this.roleMessageType,
            content: content,
            metaDict: this.metaDict
        });
    }

    extractTextAndCodePrompts(): [TextPrompt[], CodePrompt[]] {
        let textPrompts: TextPrompt[] = [];
        let codePrompts: CodePrompt[] = [];

        const lines = this.content.split("\n");
        let idx = 0;
        let startIdx = 0;

        while (idx < lines.length) {
            // 索引到代码块
            while (idx < lines.length && !lines[idx].trim().startsWith("```")) {
                idx++;
            }

            // 收集文本块
            const text = lines.slice(startIdx, idx).join("\n").trim();
            if (text) {
                textPrompts.push(new TextPrompt(text));
            }
            if (idx >= lines.length) {
                break;
            }

            // 获取代码类型
            const codeType = lines[idx].trim().substring(3).trim();
            idx++;
            startIdx = idx;

            // 索引到代码块结束
            while (!lines[idx].trim().startsWith("```")) {
                idx++;
            }

            // 收集代码块
            const code = lines.slice(startIdx, idx).join("\n").trim();
            codePrompts.push(new CodePrompt(code, codeType));
            idx++;
            startIdx = idx;
        }

        return [textPrompts, codePrompts];
    }

    getContentPrompt(): TextPrompt {
        return new TextPrompt(this.content);
    }

    getConfig(): BaseMessageConfig {
        return {
            roleName: this.roleName,
            roleType: this.roleType,
            roleMessageType: this.roleMessageType,
            content: this.content,
            metaDict: this.metaDict
        };
    }
}
